import pandas as pd
import json
class xms_ths_stock:
    def __init__(self,path=r'D:\ths\同花顺\账户\custom_block'):
        '''
        索普同花顺自选股操作,实时数据变动
        path同花顺自定义自选股路径
        支持动态板块,支持服务器
        作者:索普量化
        微信:xms_quants1
        '''
        print('''
            作者:索普量化
            微信:xms_quants1
            量化网页:http://14.103.193.242:9999/xms_quants.html
            ''')
        self.path=path
    def adjust_stock(self,stock='600031.SH'):
        '''
        调整代码
        '''
        if stock[-2:]=='SH' or stock[-2:]=='SZ' or stock[-2:]=='sh' or stock[-2:]=='sz':
            stock=stock.upper()
        else:
            if stock[:3] in ['600','601','603','605','688','689',
                             ] or stock[:2] in ['11','51','58']:
                stock=stock+'.SH'
            else:
                stock=stock+'.SZ'
        return stock
    def adjust_stock_maker_code(self,stock='513100'):
        '''
        调整股票市场代码
        自己测试修改
        '''
        stock=str(stock)[:6]
        if stock[:2] in ['51']:
            return '20'
        elif stock[:2] in ['15']:
            return '36'
        elif stock[:2] in ['60','68']:
            return '17'
        elif stock[:2] in ['00','30']:
            return '33'
        elif stock[:2] in ['11']:
            return '19'
        elif stock[:2] in ['12']:
            return '35'
        else:
            return '17'
    def get_read_ths_data(self,code='40'):
        '''
        读取原始数据
        '''
        try:
            path=r'{}\{}'.format(self.path,code)
            with open(path, 'r') as file:
                com = file.read()
            text=json.loads(com)
            return True,text
        except Exception  as e:
            print(e)
            return False,''
    def get_read_ths_stock(self,code='40'):
        '''
        234，35，40
        读取同花顺自选股板块成分股
        '''
        path=r'{}\{}'.format(self.path,code)
        df=pd.DataFrame()
        try:
            with open(path, 'r') as file:
                    com = file.read()
            text=json.loads(com)
            context=text['context']
            
            if len(context)>=6:
                stock_list=str(context).split(',')[0]
                stock_list=stock_list.split('|')
                df['证券代码']=stock_list
                df=df[:-1]
                df['证券代码']=df['证券代码'].apply(lambda  x: self.adjust_stock(x))
            else:
                print(ConnectionAbortedError,'板块没有数据')
        except Exception as e:
            print(e,code,'有问题')

        return df
    def get_add_stock_to_ths(self,code='40',stock_lists=['159351','159981']):
        '''
        添加股票到同花顺自选股
        仔细检查股票代码
        '''
        stats,text=self.get_read_ths_data(code=code)
        path=r'{}\{}'.format(self.path,code)
        if stats:
            context=text['context']
            stock_list=str(context).split(',')[0]
            stock_list=stock_list.split('|')[:-1]
            code_list=str(context).split(',')[-1]
            code_list=code_list.split('|')[:-1]
            for stock in stock_lists:
                if stock not in stock_list:
                    maker=self.adjust_stock_maker_code(stock=stock)
                    stock_list.append(stock)
                    code_list.append(maker)
                else:
                    print(stock ,'在同花顺{} 板块存在不添加'.format(code))
            stock_text='|'.join(stock_list)+'|'
            code_text='|'.join(code_list)+'|'
            context=stock_text+','+code_text
            #print(stock_list)
            #print(code_list)
            #print(stock_text)
            #print(code_text)
            #print(context)
            text['context']=context
            text=str(text)
            text=text.replace("'",'"')
            path=r'{}\{}'.format(self.path,code)
            with open(path, 'w+') as file:
                file.write(str(text))
            print(stock_lists,'添加到同花顺{}板块 成功'.format(code))
        else:
            print(stock_lists,'添加有问题')
    def get_del_stock_to_ths(self,code='40',stock_lists=['513500']):
        '''
        股票从同花顺自选股删除
        '''
        stats,text=self.get_read_ths_data(code=code)
        path=r'{}\{}'.format(self.path,code)
        if stats:
            context=text['context']
            if len(context)>=6:
                stock_list=str(context).split(',')[0]
                stock_list=stock_list.split('|')[:-1]
                code_list=str(context).split(',')[-1]
                code_list=code_list.split('|')[:-1]
                code_dict=dict(zip(stock_list,code_list))
                for stock in stock_lists:
                    if stock in stock_list:
                        stock_list.remove(stock)
                    else:
                        print(stock,'没有在同花顺板块{}不操作'.format(code))
                maker_list=[]
                for stock in stock_list:
                    maker=code_dict.get(stock,0)
                    if maker !=0:
                        maker_list.append(maker)
                code_text=maker_list
                stock_text='|'.join(stock_list)+'|'
                code_text='|'.join(code_list)+'|'
                context=stock_text+','+code_text
                #print(stock_list)
                #print(code_list)
                #print(stock_text)
                #print(code_text)
                #print(context)
                text['context']=context
                text=str(text)
                text=text.replace("'",'"')
                path=r'{}\{}'.format(self.path,code)
                with open(path, 'w+') as file:
                    file.write(str(text))
                print(stock_lists,'从同花顺{}板块删除 成功'.format(code))
            else:
                print(stock_lists,'删除有问题')
    def get_del_all_stock(self,code='40'):
        '''
        清空同花顺自选股板块数据
        '''
        stats,text=self.get_read_ths_data(code=code)
        path=r'{}\{}'.format(self.path,code)
        if stats:
            context=','
            text['context']=context
            text=str(text)
            text=text.replace("'",'"')
            path=r'{}\{}'.format(self.path,code)
            with open(path, 'w+') as file:
                file.write(str(text))
            print('同花顺板块{}清空数据完成'.format(code))
        else:
            print('清空板块数据有问题')
    
if __name__=="__main__":
    ths=xms_ths_stock()
    df=ths.get_read_ths_stock(code='330')
    print(df)

